# INT-4 SVDQuant
from diffusers import FluxPipeline, FluxTransformer2DModel
from diffusers.quantizers.quantization_config import SVDQuantConfig
import torch

ckpt_id   = "black-forest-labs/FLUX.1-dev"
quant_id  = "mit-han-lab/svdq-int4-flux.1-dev"

transformer = FluxTransformer2DModel.from_single_file(
    quant_id,
    quantization_config=SVDQuantConfig(),
    torch_dtype=torch.bfloat16,
    device_map="cuda",
)

pipe = FluxPipeline.from_pretrained(
    ckpt_id,
    transformer=transformer,
    torch_dtype=torch.bfloat16,
).to("cuda")

pipe_kwargs = {
    "prompt": "A cat holding a sign that says hello world",
    "height": 1024,
    "width": 1024,
    "guidance_scale": 3.5,
    "num_inference_steps": 50,
}
image = pipe(generator=torch.manual_seed(0), **pipe_kwargs).images[0]
image.save("svdq_int4.png")


# BF16 baseline
pipe = FluxPipeline.from_pretrained(
    ckpt_id,
    torch_dtype=torch.bfloat16,
).to("cuda")

image = pipe(generator=torch.manual_seed(0), **pipe_kwargs).images[0]
image.save("bf16.png")